import subprocess
import re
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import time
import concurrent.futures
import argparse

TARGET = "CreateThread"


def run_experiment(num_threads, sleep_duration, num_runs=5):
    """Runs the CreateThread program multiple times and returns a list of counts."""
    global TARGET
    counts = []
    for i in range(num_runs):
        try:
            result = subprocess.run(
                ["xmake", "run", TARGET, str(num_threads), str(sleep_duration)],
                capture_output=True,
                text=True,
                check=True,
            )
            output = result.stdout.strip()
            match = re.search(r"Final count: (\d+)", output)
            if match:
                count = int(match.group(1))
                counts.append(count)
            else:
                print(f"Error: Could not extract final count from output: {output}")
                return None
        except subprocess.CalledProcessError as e:
            print(f"Error running CreateThread: {e}")
            print(f"Stderr: {e.stderr}")
            return None
        time.sleep(0.1)
    return counts


def worker(num_threads, sleep_duration, num_runs):
    """Worker function for parallel execution."""
    print(
        f"Running experiment with {num_threads} threads and {sleep_duration} sleep duration..."
    )
    counts = run_experiment(num_threads, sleep_duration, num_runs)
    return (num_threads, sleep_duration), counts


def main():
    """Main function to run experiments and plot the results."""
    num_threads_list = [10, 50, 100, 500]
    sleep_duration_list = [1, 50, 100, 500]
    num_runs = 50
    results = {}

    # Use a ThreadPoolExecutor to run experiments in parallel
    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = []
        for num_threads in num_threads_list:
            for sleep_duration in sleep_duration_list:
                futures.append(
                    executor.submit(worker, num_threads, sleep_duration, num_runs)
                )

        # Collect the results as they become available
        for future in concurrent.futures.as_completed(futures):
            (num_threads, sleep_duration), counts = future.result()
            if counts is not None:
                results[(num_threads, sleep_duration)] = counts

    plot_results(results, num_threads_list, sleep_duration_list)


def plot_results(results, num_threads_list, sleep_duration_list):
    """Plots the results of the experiments using subplots."""
    fig, axes = plt.subplots(
        len(num_threads_list),
        len(sleep_duration_list),
        figsize=(15, 10),
        sharey=True,
        dpi=300,
    )

    # Flatten the axes array for easy indexing
    axes = axes.flatten()

    # Iterate through all combinations of num_threads and sleep_duration
    for i, num_threads in enumerate(num_threads_list):
        for j, sleep_duration in enumerate(sleep_duration_list):
            ax = axes[i * len(sleep_duration_list) + j]  # Get the current subplot

            # Prepare data for the current subplot
            data = []
            if (num_threads, sleep_duration) in results and results[
                (num_threads, sleep_duration)
            ]:
                for count in results[(num_threads, sleep_duration)]:
                    data.append({"Count": count})
            df = pd.DataFrame(data)

            # Create the boxplot
            sns.histplot(df, x="Count", bins=20, ax=ax)

            # Set the title for the subplot
            ax.set_title(f"Threads: {num_threads}, Sleep: {sleep_duration}")
            ax.set_xlabel("")  # Remove x label
            if j == 0:
                ax.set_ylabel("Final Count")  # Y label only on the first column
            else:
                ax.set_ylabel("")  # Remove y label for other columns

    plt.tight_layout()
    plt.savefig("results_subplots.png")
    plt.show()


if __name__ == "__main__":
    args = argparse.ArgumentParser()
    args.add_argument("--target", type=str)
    args = args.parse_args()
    TARGET = args.target
    main()
